#!/usr/bin/ruby

# Summary: monitor who is eating your bandwidth
# WWW:     http://code.google.com/p/trafmonitor/
# Author:  tmarciniak@gmail.com
# License: New BSD

# WISHLIST
# - make detailed statistics "per user+protocol" and "global+protocol"
# - gather data into SQL db (sqlite?)
# - realtime monitoring (eg curses based, top 10 downloaders)
# - bidirectional monitoring?

require 'pcaplet'

def dump_stats( stats, logfile_path )
    File.open( logfile_path, 'a' ) { |fh|
        fh.puts 'Traffic summary dump at: ' + Time.now.to_s
        stats.each { |key, value|
            fh.puts key.to_s + ' ' + value.to_s
        }
    }
end

def pid_handler( pidfile_path )
    if  File.exist?( pidfile_path )
        puts "Is trafmonitor already running? Check #{pidfile_path}. Terminating."
        exit 1
    else
        File.open( pidfile_path, 'w+' ) { |fh|
            fh.puts Process.pid
        }
    end
end

def main
    logfile_path = '/var/log/trafmonitor.log'
    pidfile_path = '/var/run/trafmonitor.pid'

    filterexp = 'tcp and dst net 192.168.0.0/16 and src net !192.168.0.0/16'
    stats     = {}
    sniffer   = Pcaplet.new( '-i eth0 -s 1500' )
    my_filter = Pcap::Filter.new( filterexp, sniffer.capture )
    
    pid_handler( pidfile_path )
    
    sniffer.add_filter( my_filter )
    
    sniffer.each_packet { |p|
        # we check p.ip_dst.to_num_s != nil; this may seem unclear at first sight - what 
        # we do here is that we check if stats[ p.ip_dst.to_num_s ] is populated by key 
        # so it can be further manipulated by value; just a ruby quirk
        stats[ p.ip_dst.to_num_s ] = stats.fetch( p.ip_dst.to_num_s, 0 ) + p.size.to_i 
    
        Signal.trap( "USR1" ) do
            dump_stats( stats, logfile_path )
        end
    
        Signal.trap( "USR2" ) do
            File.open( logfile_path, 'a' ) { |fh|
                fh.puts 'Stats cleared with SIGUSR2 at: ' + Time.now.to_s
            }

            stats = Hash.new
        end

        Signal.trap( "INT" ) do
            File.open( logfile_path, 'a' ) { |fh|
                fh.puts "Caught interrupt... Dumping statistics and terminating."
            }
            
            dump_stats( stats, logfile_path )
            File.delete( pidfile_path )
            exit 0
        end
    }
end

pid = fork do
    main()
end

puts "trafmonitor daemon started with pid #{pid}."

Process.detach( pid )
